import numpy as np
import matplotlib.pyplot as plt
from python_ai.ML.none_lin.square.the_data import *
from python_ai.ML.lin_regression.xlib import *


def xget_linspace(xmin, xmax, mul=10):
    return np.linspace(xmin, xmax, int(np.ceil(xmax-xmin) * mul))


np.random.seed(1)
xmin = -100
xmax = 1200
a, b, c = 0.05, -20, -10000
m = 20
delta = 400
iter_ori = iter = 15000
spr = 1
spc = 3
spn = 0
plt.figure(figsize=[16, 5])

# original data and plot
spn += 1
plt1 = plt.subplot(spr, spc, spn)
plt_x, plt_y = get_square_data(xmin, xmax, m, [a, b, c], delta)
plt_x_min = plt_x.min()
plt_x_max = plt_x.max()
plt1.scatter(plt_x, plt_y, s=5)

# map to multi column
x0 = np.ones([m, 1])
x1 = plt_x.reshape(m, 1)
x2 = (plt_x ** 2).reshape(m, 1)
x12 = np.c_[x1, x2]
y = plt_y.reshape(m, 1)
plt_xy = np.c_[x12, y]

# scale mapped data and plot
spn += 1
plt2 = plt.subplot(spr, spc, spn)
# scale mapped data
plt_xy_scaled, mu_vec, sigma_vec = scale_feature_data(plt_xy)
# plot dots
plt_x_scaled = plt_xy_scaled[:, 0]
plt_x_scaled_min = plt_x_scaled.min()
plt_x_scaled_max = plt_x_scaled.max()
plt_y_scaled = plt_xy_scaled[:, -1]
plt2.scatter(plt_x_scaled, plt_y_scaled, s=5)
# mapped data
x1_scaled = plt_xy_scaled[:, 0].reshape(m, 1)
x2_scaled = plt_xy_scaled[:, 1].reshape(m, 1)
y_scaled = plt_y_scaled.reshape(m, 1)
XX_scaled = np.c_[x0, x1_scaled, x2_scaled]

# gradient descent and plot
theta_scaled, j_history, xscores = gradient_descent_algorithm(XX_scaled, y_scaled, num_iters=iter)
if iter_ori > iter:
    print(f'Converged at {iter}th iteration!')
else:
    print(f'Not converged after {iter} iterations!!!')
print(f'Theta = {theta_scaled}, final J = {j_history[-1]}, final score = {xscores[-1]}')
plt_func_x_scaled = xget_linspace(plt_x_scaled_min, plt_x_scaled_max)
plt_func_y_scaled = square_func(plt_func_x_scaled, theta_scaled[2], theta_scaled[1], theta_scaled[0])
plt2.plot(plt_func_x_scaled, plt_func_y_scaled, 'r-')

# unscale and plot
theta = scale_theta_back(x12, y, theta_scaled)
print(f'Theta unscaled = {theta}')
plt_func_x = xget_linspace(plt_x_min, plt_x_max)
plt_func_y = square_func(plt_func_x, theta[2], theta[1], theta[0])
plt1.plot(plt_func_x, plt_func_y, 'r-')
plt_func_tgt = square_func(plt_func_x, a, b, c)
plt1.plot(plt_func_x, plt_func_tgt, 'g--')

spn += 1
plt3 = plt.subplot(spr, spc, spn)
plt3.plot(j_history)

plt1.grid()
plt2.grid()
plt3.grid()
plt.show()
